package io.github.yotta.common.utils.tree;

import lombok.Builder;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * 树形结构构建工具类
 *
 * @param <T> 节点标识类型（如 id 和 parentId）
 * @param <D> 用户传入的原始数据类型
 */
@Builder
public class TreeBuilder<T, D> {
    // 提取节点 ID 的函数
    private Function<D, T> idSelector;
    // 提取父节点 ID 的函数
    private Function<D, T> parentIdSelector;
    // 创建 TreeNode 的工厂方法
    private Function<D, TreeNode<D>> nodeFactory = TreeNode::new;
    // 判断是否为根节点的函数
    private Predicate<T> isRoot = new DefaultIsRoot<T>();

    public static class DefaultIsRoot<T> implements Predicate<T> {
        @Override
        public boolean test(T t) {
            if (t == null) {
                return true;
            }
            String s = String.valueOf(t);
            return s.isEmpty() || s.equals("0");
        }
    }

    /**
     * 构建树形结构
     *
     * @param dataList 平铺的用户数据列表，每个元素为用户定义的泛型类型 D
     * @return 根节点列表，每个根节点为 TreeNode 类型
     * 
     * 使用示例：
     * 1. 定义 idSelector：从用户数据中提取唯一标识，例如 (data) -> data.getId()
     * 2. 定义 parentIdSelector：从用户数据中提取父节点标识，例如 (data) -> data.getParentId()
     * 3. 定义 nodeFactory：将用户数据转换为 TreeNode 实例，例如 (data) -> new TreeNode<>()
     * 4. 调用 buildTree 方法传入用户数据列表即可生成树形结构
     */
    public List<TreeNode<D>> buildTree(List<D> dataList) {
        if (dataList == null || dataList.isEmpty()) {
            return Collections.emptyList();
        }

        // 节点映射表，用于快速查找节点
        Map<T, TreeNode<D>> nodeMap = dataList.stream()
                .map(data -> {
                    TreeNode<D> node = nodeFactory.apply(data);
                    node.setData(data);
                    return node;
                })
                .collect(Collectors.toMap(data -> idSelector.apply(data.getData()),
                        node -> node));

        // 根节点列表
        List<TreeNode<D>> rootNodes = new ArrayList<>();

        for (TreeNode<D> node : nodeMap.values()) {
            T parentId = parentIdSelector.apply(node.getData());
            if (isRoot.test(parentId) || !nodeMap.containsKey(parentId)) {
                // 判断父节点是不是根节点，比如说根节点的 parentId 为 null 或不存在
                rootNodes.add(node);
            } else {
                // 否则将当前节点添加到父节点的子节点列表中
                TreeNode<D> parentNode = nodeMap.get(parentId);
                parentNode.addChild(node);
            }
        }

        return rootNodes;
    }
}